package commons.base.http;

import java.io.FileInputStream;
import java.io.IOException;
import java.net.UnknownHostException;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;

import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLServerSocket;
import javax.net.ssl.SSLSession;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

/**
 * 简化构建SSLContext操作
 * @author yuan<cihang.yuan@happyelements.com>
 *
 */
public class SSLBuilder {

	/**
	 * HostnameVerifier的一个简单实现, 允许所有hostname.
	 */
	private HostnameVerifier hostnameVerifier = new HostnameVerifier() {
		@Override
		public boolean verify(String hostname, SSLSession session) {
			return true;
		}
	};
	
	/** 秘钥管理器 */
	private KeyManager[] kms;
	/** 信任管理器, 用来管理证书, 验证对方 */
	private TrustManager[] tms;
	/** 安全随机数 */
	private SecureRandom secureRandom;
	
	/**
	 * 加载JKS类型(JDK默认类型)文件存储的秘钥
	 * @param file
	 * @param password 密码
	 * @return
	 */
	public SSLBuilder loadDefaultTypeKey(String file, String password){
		// jks
		loadKey(KeyStore.getDefaultType(), file, password);
		return this;
	}
	
	/**
	 * 加载P12类型文件存储的秘钥
	 * @param file
	 * @param password 密码
	 * @return
	 */
	public SSLBuilder loadP12Key(String file, String password){
		loadKey("PKCS12", file, password);
		return this;
	}
	
	/**
	 * 从文件加载存储的秘钥
	 * @param type 文件类型
	 * @param file 文件路径
	 * @param password 密码
	 * @return
	 */
	public SSLBuilder loadKey(String type, String file, String password){
		FileInputStream fis = null;
		
		try{
			fis = new FileInputStream(file);
			
			KeyStore clientStore = KeyStore.getInstance(type);
			clientStore.load(fis, password.toCharArray());
			
			KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
			kmf.init(clientStore, password.toCharArray());
			
			this.kms = kmf.getKeyManagers();
		}catch(Exception e){
			throw new RuntimeException(e.getMessage(), e);
		}finally{
			try {
				if(fis != null) fis.close();
			} catch (IOException e) {
				throw new RuntimeException(e.getMessage(), e);
			}
		}
		
		return this;
	}
	
	/**
	 * 加载JKS类型(JDK默认类型)文件存储的信任证书. 与useAllowAllTrust方法互斥
	 * @param file 文件路径
	 * @param password 密码
	 * @return
	 */
	public SSLBuilder loadDefaultTypeTrust(String file, String password){
		// jks
		loadTrust(KeyStore.getDefaultType(), file, password);
		
		return this;
	}
	
	/**
	 * 加载P12类型文件存储的信任证书. 与useAllowAllTrust方法互斥
	 * @param file 文件路径
	 * @param password 密码
	 * @return
	 */
	public SSLBuilder loadP12Trust(String file, String password){
		loadTrust("PKCS12", file, password);
		return this;
	}
	
	/**
	 * 从文件加载存储的信任证书. 与useAllowAllTrust方法互斥
	 * @param type 文件类型
	 * @param file 文件路径
	 * @param password 密码
	 * @return
	 */
	public SSLBuilder loadTrust(String type, String file, String password) {
		FileInputStream fis = null;
		
		try{
			fis = new FileInputStream(file);
			
			KeyStore trustStore = KeyStore.getInstance(type);
			trustStore.load(fis, password.toCharArray());
			
			TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
			tmf.init(trustStore);
			
			this.tms = tmf.getTrustManagers();
		}catch(Exception e){
			throw new RuntimeException(e.getMessage(), e);
		}finally{
			try {
				if(fis != null) fis.close();
			} catch (IOException e) {
				throw new RuntimeException(e.getMessage(), e);
			}
		}
		return this;
	}
	
	/**
	 * 信任所有的客户端和服务器 . 与 loadTrust, loadDefaultTypeTrust,loadP12Trust方法互斥
	 * @return
	 */
	public SSLBuilder useAllowAllTrust() {
		X509TrustManager tm = new X509TrustManager() {
			public java.security.cert.X509Certificate[] getAcceptedIssuers() {
				return null;
			}

			public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException {
			}

			public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateException {
			}
		};

		this.tms = new TrustManager[] { tm };
		
		return this;
	}
	
	/**
	 * 使用安全随机数
	 * @return
	 */
	public SSLBuilder useSecureRandom(){
		this.secureRandom = new SecureRandom();
		
		return this;
	}
	
	/**
	 * 自定义HostnameVerifier, 默认允许所有hostname
	 * @param hostnameVerifier
	 * @return
	 */
	public SSLBuilder setHostnameVerifier(HostnameVerifier hostnameVerifier){
		this.hostnameVerifier = hostnameVerifier;
		return this;
	}
	
	/**
	 * 构造SSLContext
	 * @return
	 */
	public SSLContext build(){
		try {
			SSLContext ctx = SSLContext.getInstance("TLS");
			
			ctx.init(kms, tms, secureRandom);
			
			return ctx;
		}catch(Exception e) {
			throw new RuntimeException(e.getMessage(), e);
		}
	}
	
	/**
	 * 修改JDK默认https设置
	 */
	public void setJDKDefaultHttps(){
		SSLContext ssl = build();
		
		HttpsURLConnection.setDefaultSSLSocketFactory(ssl.getSocketFactory());
		HttpsURLConnection.setDefaultHostnameVerifier(hostnameVerifier);
	}
	
	/**
	 * 构造SSLServerSocket
	 * @param port 端口
	 * @return
	 * @throws IOException
	 */
	public SSLServerSocket createSSLServerSocket(int port) throws IOException{
		return (SSLServerSocket)build().getServerSocketFactory().createServerSocket(port);
	}
	
	/**
	 * 构造SSLSocket
	 * @param host IP
	 * @param port 端口
	 * @return
	 * @throws UnknownHostException
	 * @throws IOException
	 */
	public SSLSocket createSSLSocket(String host, int port) throws UnknownHostException, IOException{
		return (SSLSocket)build().getSocketFactory().createSocket(host, port);
	}
	
}
